package com.als;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class SimpleALS {
    private final int numUsers;
    private final int numItems;
    private final int rank;
    private final double[][] userFeatures;
    private final double[][] itemFeatures;
    private final double[][] ratings;

    public SimpleALS(int numUsers, int numItems, int rank) {
        this.numUsers = numUsers;
        this.numItems = numItems;
        this.rank = rank;
        this.userFeatures = new double[numUsers][rank];
        this.itemFeatures = new double[numItems][rank];
        this.ratings = new double[numUsers][numItems];

        // 初始化特征矩阵
        Random rand = new Random(42);
        for (int i = 0; i < numUsers; i++) {
            for (int j = 0; j < rank; j++) {
                userFeatures[i][j] = rand.nextDouble();
            }
        }
        for (int i = 0; i < numItems; i++) {
            for (int j = 0; j < rank; j++) {
                itemFeatures[i][j] = rand.nextDouble();
            }
        }
    }

    public static void main(String[] args) {
        SimpleALS als = new SimpleALS(5, 10, 3);

        // 设置一些示例评分
        als.setRating(0, 0, 5.0);
        als.setRating(0, 1, 3.0);
        als.setRating(1, 0, 4.0);
        als.setRating(1, 2, 2.0);

        // 训练模型
        als.train(10, 0.01);

        // 生成推荐
        List<Integer> recommendations = als.recommendForUser(0, 3);
        System.out.println("为用户 0 推荐: " + recommendations);
    }

    public void setRating(int userId, int itemId, double rating) {
        ratings[userId][itemId] = rating;
    }

    public void train(int maxIterations, double lambda) {
        for (int iter = 0; iter < maxIterations; iter++) {
            // 固定物品特征，更新用户特征
            for (int u = 0; u < numUsers; u++) {
                updateUserFeatures(u, lambda);
            }

            // 固定用户特征，更新物品特征
            for (int i = 0; i < numItems; i++) {
                updateItemFeatures(i, lambda);
            }

            double loss = calculateLoss(lambda);
            System.out.println("迭代 " + iter + ", 损失: " + loss);
        }
    }

    private void updateUserFeatures(int userId, double lambda) {
        // 简化版的用户特征更新
        // 实际实现需要解线性方程组
        for (int f = 0; f < rank; f++) {
            double sum = 0.0;
            int count = 0;

            for (int i = 0; i < numItems; i++) {
                if (ratings[userId][i] > 0) {
                    double prediction = 0.0;
                    for (int k = 0; k < rank; k++) {
                        if (k != f) {
                            prediction += userFeatures[userId][k] * itemFeatures[i][k];
                        }
                    }
                    double error = ratings[userId][i] - prediction;
                    sum += error * itemFeatures[i][f];
                    count++;
                }
            }

            if (count > 0) {
                userFeatures[userId][f] = sum / (count + lambda);
            }
        }
    }

    private void updateItemFeatures(int itemId, double lambda) {
        // 类似的物品特征更新逻辑
        for (int f = 0; f < rank; f++) {
            double sum = 0.0;
            int count = 0;

            for (int u = 0; u < numUsers; u++) {
                if (ratings[u][itemId] > 0) {
                    double prediction = 0.0;
                    for (int k = 0; k < rank; k++) {
                        if (k != f) {
                            prediction += userFeatures[u][k] * itemFeatures[itemId][k];
                        }
                    }
                    double error = ratings[u][itemId] - prediction;
                    sum += error * userFeatures[u][f];
                    count++;
                }
            }

            if (count > 0) {
                itemFeatures[itemId][f] = sum / (count + lambda);
            }
        }
    }

    private double calculateLoss(double lambda) {
        double loss = 0.0;
        int count = 0;

        for (int u = 0; u < numUsers; u++) {
            for (int i = 0; i < numItems; i++) {
                if (ratings[u][i] > 0) {
                    double prediction = predict(u, i);
                    loss += Math.pow(ratings[u][i] - prediction, 2);
                    count++;
                }
            }
        }

        // 添加正则化项
        for (int u = 0; u < numUsers; u++) {
            for (int f = 0; f < rank; f++) {
                loss += lambda * Math.pow(userFeatures[u][f], 2);
            }
        }
        for (int i = 0; i < numItems; i++) {
            for (int f = 0; f < rank; f++) {
                loss += lambda * Math.pow(itemFeatures[i][f], 2);
            }
        }

        return loss / count;
    }

    public double predict(int userId, int itemId) {
        double prediction = 0.0;
        for (int f = 0; f < rank; f++) {
            prediction += userFeatures[userId][f] * itemFeatures[itemId][f];
        }
        return prediction;
    }

    public List<Integer> recommendForUser(int userId, int topN) {
        List<Recommendation> recommendations = new ArrayList<>();

        for (int i = 0; i < numItems; i++) {
            if (ratings[userId][i] == 0) { // 只推荐未评分的物品
                double score = predict(userId, i);
                recommendations.add(new Recommendation(i, score));
            }
        }

        // 按评分排序
        recommendations.sort((a, b) -> Double.compare(b.score, a.score));

        List<Integer> result = new ArrayList<>();
        for (int i = 0; i < Math.min(topN, recommendations.size()); i++) {
            result.add(recommendations.get(i).itemId);
        }

        return result;
    }

    private static class Recommendation {
        int itemId;
        double score;

        Recommendation(int itemId, double score) {
            this.itemId = itemId;
            this.score = score;
        }
    }
}
